import cv2
import numpy as np

# 加载原始图像
original_image = cv2.imread('original_image.jpg', cv2.IMREAD_GRAYSCALE)

# 加载带有水印的图像
watermarked_image = cv2.imread('watermarked_image.jpg', cv2.IMREAD_GRAYSCALE)

# 初始化攻击强度参数
jpeg_quality = 90  # JPEG压缩质量
rotation_angle = 30  # 旋转角度

# 模拟JPEG压缩攻击
compressed_image = cv2.imencode('.jpg', watermarked_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])[1]
decompressed_image = cv2.imdecode(compressed_image, cv2.IMREAD_GRAYSCALE)

# 模拟旋转攻击
rows, cols = watermarked_image.shape
M = cv2.getRotationMatrix2D((cols / 2, rows / 2), rotation_angle, 1)
rotated_image = cv2.warpAffine(watermarked_image, M, (cols, rows))



# 计算均方误差（MSE）
# 1 通过原始图像和嵌入水印后的图像对比，比较嵌入水印后的图像和原始图像之间的差异，越小越好
# 2 通过原始图像和嵌入水印后的图像对比，评估嵌入水印后的图像质量，失真越多MSE越大吗，越小越好
# 3 嵌入水印后的图像和受攻击后的带水印图像对比，受攻击后图像可提取水印的能力即算法鲁棒性，越小越好
mse_compression = np.mean((decompressed_image - original_image) ** 2)
mse_rotation = np.mean((rotated_image - original_image) ** 2)

# 打印结果
print(f"JPEG Compression Attack MSE: {mse_compression}")
print(f"Rotation Attack MSE: {mse_rotation}")

# 计算信噪比 (SNR)
# 1 评估嵌入水印后图像质量，越高越好
# 2 评估受攻击后的图像提取水印的难易度，越高越好
# 3 评估嵌入图像的水印的强度，越高水印强度越高
# 4 评估提取水印的质量，越高越好
snr_original = 10 * np.log10(np.mean(original_image ** 2) / np.mean((original_image - watermarked_image) ** 2))
snr_compression = 10 * np.log10(np.mean(original_image ** 2) / np.mean((original_image - decompressed_image) ** 2))
snr_rotation = 10 * np.log10(np.mean(original_image ** 2) / np.mean((original_image - rotated_image) ** 2))

# 打印结果
print(f"SNR of Original Image vs. Watermarked Image: {snr_original} dB")
print(f"SNR after JPEG Compression Attack: {snr_compression} dB")
print(f"SNR after Rotation Attack: {snr_rotation} dB")



# 加载原始水印图像
original_watermark = cv2.imread('original_watermark.png', cv2.IMREAD_GRAYSCALE)

# 加载提取的水印图像
extracted_watermark = cv2.imread('extracted_watermark.png', cv2.IMREAD_GRAYSCALE)

# 计算相关性
correlation = np.corrcoef(original_watermark.flatten(), extracted_watermark.flatten())[0, 1]

# 打印结果
print(f"Correlation between Original and Extracted Watermark: {correlation}")